{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 完全从零写一个 MOE 大模型（LLM）\n",
    "## Build a miniMoE-LLM from scratch\n",
    "- 从零开始构建 MoE 模型，从基础讲起，共讲解三个版本\n",
    "   - 基础版本，理解MOE\n",
    "   - SparseMoE，了解大模型怎么训练 MOE LLM\n",
    "   - ShareExpert SparseMoe，了解 deepseek 训练 MOE 模型算法\n",
    "- 运行代码完全手写，一行一行带着运行\n",
    "- 作者： **[chaofa用代码打点酱油](https://bruceyuan.com/about.html)**\n",
    "\n",
    "欢迎阅读原文：[LLM MOE的进化之路，从普通简化 MOE，到 spare_moe，再到 deepseek 使用的 share_xpert_spare_moe](https://bruceyuan.com/llms-zero-to-hero/the-way-of-moe-model-evolution.html)\n",
    "\n",
    "原文中有对应的图片，和简单的讲解，可以方便你理解\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "## 1.  版本1：基础版本MOE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BasicExpert(nn.Module):\n",
    "    # 一个 Expert 可以是一个最简单的， linear 层即可\n",
    "    # 也可以是 MLP 层\n",
    "    # 也可以是 更复杂的 MLP 层（active function 设置为 swiglu）\n",
    "    def __init__(self, feature_in, feature_out):\n",
    "        super().__init__()\n",
    "        self.linear = nn.Linear(feature_in, feature_out)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return self.linear(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class BasicMOE(nn.Module):\n",
    "    def __init__(self, feature_in, feature_out, expert_number):\n",
    "        super().__init__()\n",
    "        self.experts = nn.ModuleList(\n",
    "            [\n",
    "                BasicExpert(feature_in, feature_out) for _ in range(expert_number)\n",
    "            ]\n",
    "        )\n",
    "        # gate 就是选一个 expert \n",
    "        self.gate = nn.Linear(feature_in, expert_number)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        # x 的 shape 是 （batch, feature_in)\n",
    "        expert_weight = self.gate(x)  # shape 是 (batch, expert_number)\n",
    "        expert_out_list = [\n",
    "            expert(x).unsqueeze(1) for expert in self.experts\n",
    "        ]  # 里面每一个元素的 shape 是： (batch, ) ??\n",
    "\n",
    "        # concat 起来 (batch, expert_number, feature_out)\n",
    "        expert_output = torch.cat(expert_out_list, dim=1)\n",
    "\n",
    "        # print(expert_output.size())\n",
    "\n",
    "        expert_weight = expert_weight.unsqueeze(1) # (batch, 1, expert_nuber)\n",
    "\n",
    "        # expert_weight * expert_out_list\n",
    "        output = expert_weight @ expert_output  # (batch, 1, feature_out)\n",
    "        \n",
    "        return output.squeeze()\n",
    "\n",
    "\n",
    "def test_basic_moe():\n",
    "    x = torch.rand(2, 4)\n",
    "\n",
    "    basic_moe = BasicMOE(4, 3, 2)\n",
    "    out = basic_moe(x)\n",
    "    print(out)\n",
    "\n",
    "\n",
    "test_basic_moe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 版本2：SparseMoE （大模型训练使用）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# 主要参考自 mistral MOE 的实现\n",
    "\n",
    "class MOERouter(nn.Module):\n",
    "    def __init__(self, hidden_dim, expert_number, top_k):\n",
    "        super().__init__()\n",
    "        self.gate = nn.Linear(hidden_dim, expert_number)\n",
    "        self.expert_number = expert_number\n",
    "        self.top_k = top_k\n",
    "    \n",
    "    def forward(self, hidden_states):\n",
    "        # 计算路由logits\n",
    "        router_logits = self.gate(hidden_states)  # shape is (b * s, expert_number)\n",
    "        \n",
    "        # 计算专家经过softmax之后的概率\n",
    "        routing_probs = F.softmax(router_logits, dim=-1, dtype=torch.float)\n",
    "        \n",
    "        # 计算topk的专家的输出\n",
    "        router_weights, selected_experts = torch.topk(\n",
    "            routing_probs, self.top_k, dim=-1\n",
    "        )  # shape都是 (b * s, top_k)\n",
    "        \n",
    "        # 专家权重归一化\n",
    "        router_weights = router_weights / router_weights.sum(dim=-1, keepdim=True)\n",
    "        router_weights = router_weights.to(hidden_states.dtype)\n",
    "        \n",
    "        # 生成专家掩码\n",
    "        expert_mask = F.one_hot(\n",
    "            selected_experts,\n",
    "            num_classes=self.expert_number\n",
    "        )  # shape是 (b * s, top_k, expert_number)\n",
    "        expert_mask = expert_mask.permute(2, 1, 0)  # (expert_number, top_k, b * s)\n",
    "        \n",
    "        return router_logits, router_weights, selected_experts, expert_mask\n",
    "\n",
    "\n",
    "class MOEConfig:\n",
    "    def __init__(\n",
    "            self, \n",
    "            hidden_dim, \n",
    "            expert_number, \n",
    "            top_k, \n",
    "            shared_experts_number=2,\n",
    "        ):\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.expert_number = expert_number\n",
    "        self.top_k = top_k\n",
    "        self.shared_experts_number = shared_experts_number\n",
    "\n",
    "class SparseMOE(nn.Module):\n",
    "    # 稀疏 MOE 模型，这里每一个 token 都会过 topk 个专家，得到对应token 的 hidden_embeddings\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "\n",
    "        self.hidden_dim = config.hidden_dim\n",
    "\n",
    "        self.expert_number = config.expert_number\n",
    "        self.top_k = config.top_k\n",
    "\n",
    "        self.experts = nn.ModuleList(\n",
    "            [\n",
    "                BasicExpert(self.hidden_dim, self.hidden_dim) for _ in range(self.expert_number)\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        self.router = MOERouter(self.hidden_dim, self.expert_number, self.top_k)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        # x shape is (b, s, hidden_dim)\n",
    "        batch_size, seq_len, hidden_dim = x.size()\n",
    "\n",
    "        # 合并前两个维度，因为不是 Sample 维度了，而是 token 维度\n",
    "        hidden_states = x.view(-1, hidden_dim) # shape is(b * s, hidden_dim)\n",
    "\n",
    "        router_logits, router_weights, selected_experts_indices, expert_mask = self.router(hidden_states)\n",
    "        # 其中 selected_experts_indices shape 是 (b * s, top_k)\n",
    "        # 其中 expert_mask shape 是 (expert_number, top_k, b * s)\n",
    "        \n",
    "        final_hidden_states = torch.zeros(\n",
    "            (batch_size * seq_len, hidden_dim),\n",
    "            dtype=hidden_states.dtype,\n",
    "            device=hidden_states.device\n",
    "        )\n",
    "\n",
    "        for expert_idx in range(self.expert_number):\n",
    "            expert_layer = self.experts[expert_idx]\n",
    "            # expert_mask[expert_idx] shape 是 (top_k, b * s)\n",
    "            idx, top_x = torch.where(expert_mask[expert_idx]) \n",
    "            # idx 和 top_x 都是一维 tensor\n",
    "            # idx 的值是 0 或 1, 表示这个 token 是作为当前专家的 top1 还是 top2\n",
    "            # top_x 的值是 token 在 batch*seq_len 中的位置索引\n",
    "            # 例如对于 batch_size=2, seq_len=4 的输入:\n",
    "            # top_x 的值范围是 0-7, 表示在展平后的 8 个 token 中的位置\n",
    "            # idx 的值是 0/1, 表示这个 token 把当前专家作为其 top1/top2 专家\n",
    "\n",
    "            # hidden_states 的 shape 是 (b * s, hidden_dim)\n",
    "            # 需要取到 top_x 对应的 hidden_states\n",
    "            current_state = hidden_states.unsqueeze(\n",
    "                0\n",
    "            )[:, top_x, :].reshape(-1, hidden_dim) # （selected_token_number, hidden_dim）\n",
    "\n",
    "            # router_weight 的 shape 是 (b * s, top_k)\n",
    "            current_hidden_states = expert_layer(\n",
    "                current_state\n",
    "            ) * router_weights[top_x, idx].unsqueeze(-1)  # （selected_token_number, 1） 这里有广播\n",
    "\n",
    "            # 把当前专家的输出加到 final_hidden_states 中\n",
    "            # 方式1 的写法性能更好，并且方式1容易出现\n",
    "            final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))\n",
    "            # 方式2\n",
    "            # final_hidden_states[top_x] += current_hidden_states.to(hidden_states.dtype)\n",
    "            # 方式1 的写法性能更差，并且方式1容易出现错误，+= 操作在处理重复索引时需要多次读写内存，可能会导致竞争条件\n",
    "\n",
    "        # 把 final_hidden_states 还原到原来的 shape\n",
    "        final_hidden_states = final_hidden_states.reshape(batch_size, seq_len, hidden_dim)\n",
    "\n",
    "        return final_hidden_states, router_logits # shape 是 (b * s, expert_number)\n",
    "\n",
    "\n",
    "def test_token_level_moe():\n",
    "    x = torch.rand(2, 4, 16)\n",
    "    config = MOEConfig(16, 2, 2)\n",
    "    token_level_moe = SparseMOE(config)\n",
    "    out = token_level_moe(x)\n",
    "    print(out[0].shape, out[1].shape)\n",
    "\n",
    "\n",
    "test_token_level_moe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 版本3：ShareExpert SparseMoE （deepseek 版本）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ShareExpertMOE(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "\n",
    "        self.moe_model = SparseMOE(config)\n",
    "        self.shared_experts = nn.ModuleList(\n",
    "            [\n",
    "                BasicExpert(\n",
    "                    config.hidden_dim, config.hidden_dim\n",
    "                ) for _ in range(config.shared_experts_number)\n",
    "            ]\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        # x shape 是 (b, s, hidden_dim)\n",
    "        # 首先过 moe 模型\n",
    "        sparse_moe_out, router_logits = self.moe_model(x)\n",
    "        \n",
    "        # 针对的还是 x 的每一个 \n",
    "        # 然后过 shared experts\n",
    "        shared_experts_out = [\n",
    "            expert(x) for expert in self.shared_experts\n",
    "        ] # 每一个 expert 的输出 shape 是 (b, s, hidden_dim)\n",
    "        \n",
    "        shared_experts_out = torch.stack(\n",
    "            shared_experts_out, dim=0\n",
    "        ).sum(dim=0)\n",
    "        \n",
    "        # 把 sparse_moe_out 和 shared_experts_out 加起来\n",
    "        return sparse_moe_out + shared_experts_out, router_logits\n",
    "\n",
    "\n",
    "def test_share_expert_moe():\n",
    "    x = torch.rand(2, 4, 16)\n",
    "    config = MOEConfig(16, 2, 2)\n",
    "    share_expert_moe = ShareExpertMOE(config)\n",
    "    out = share_expert_moe(x)\n",
    "    print(out[0].shape, out[1].shape)\n",
    "\n",
    "\n",
    "test_share_expert_moe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 测试， loss 部分为 deepseek 生成；\n",
    "\n",
    "def switch_load_balancing_loss(router_logits: torch.Tensor, num_experts: int) -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    计算 Switch Transformers 的负载均衡损失\n",
    "    \n",
    "    Args:\n",
    "        router_logits: shape [batch_size * sequence_length, num_experts]\n",
    "        num_experts: 专家数量\n",
    "    \n",
    "    Returns:\n",
    "        total_loss: 总损失 = auxiliary_loss + z_loss\n",
    "    \"\"\"\n",
    "    # 计算路由概率\n",
    "    router_probs = torch.softmax(router_logits, dim=-1)  # [b*s, num_experts]\n",
    "    \n",
    "    # 获取每个token的最优专家\n",
    "    _, selected_experts = torch.topk(router_probs, k=2, dim=-1)  # [b*s]\n",
    "    \n",
    "    # 创建one-hot矩阵表示选中的专家\n",
    "    mask = torch.nn.functional.one_hot(selected_experts, num_experts).float()  # [b*s, num_experts]\n",
    "    \n",
    "    # 计算每个专家的期望负载 (理想情况下应该是 1/num_experts)\n",
    "    expected_load = torch.ones_like(router_probs) / num_experts\n",
    "    \n",
    "    # 计算实际负载 (每个专家处理的token数量除以总token数量)\n",
    "    # 在batch维度上计算平均值\n",
    "    actual_load = mask.mean(dim=0)  # [num_experts]\n",
    "    \n",
    "    # 计算auxiliary loss\n",
    "    # 这会惩罚负载分布与期望负载的差异\n",
    "    aux_loss = torch.sum(actual_load * router_probs.mean(dim=0)) * num_experts\n",
    "    \n",
    "    # 计算z_loss (可选)\n",
    "    # 这会惩罚过大的路由logits\n",
    "    z_loss = torch.mean(torch.square(router_logits))\n",
    "    z_loss_weight = 0.001  # 可调整的超参数\n",
    "    \n",
    "    # 总损失\n",
    "    total_loss = aux_loss + z_loss * z_loss_weight\n",
    "    \n",
    "    return total_loss\n",
    "\n",
    "def test_moe_training():\n",
    "    # Create a simple dataset\n",
    "    batch_size = 32\n",
    "    seq_len = 16\n",
    "    hidden_dim = 32\n",
    "    num_batches = 100\n",
    "    \n",
    "    # Initialize model and optimizer\n",
    "    config = MOEConfig(hidden_dim=hidden_dim, \n",
    "                      expert_number=4,\n",
    "                      top_k=2,\n",
    "                      shared_experts_number=2)\n",
    "    model = ShareExpertMOE(config)\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "    \n",
    "    # Training loop\n",
    "    model.train()\n",
    "    for batch in range(num_batches):\n",
    "        # Generate random input data\n",
    "        x = torch.randn(batch_size, seq_len, hidden_dim)\n",
    "        target = torch.randn(batch_size, seq_len, hidden_dim)\n",
    "        \n",
    "        # Forward pass\n",
    "        output, router_logits = model(x)\n",
    "\n",
    "        # Compute losses\n",
    "        # MSE loss for prediction\n",
    "        mse_loss = F.mse_loss(output, target)\n",
    "        \n",
    "        aux_loss = switch_load_balancing_loss(router_logits, config.expert_number)\n",
    "        # Combined loss\n",
    "        total_loss = mse_loss + 0.01 * aux_loss\n",
    "        \n",
    "        # Backward pass and optimize\n",
    "        optimizer.zero_grad()\n",
    "        total_loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        if batch % 10 == 0:\n",
    "            print(f\"Batch {batch}, Loss: {total_loss.item():.4f} \"\n",
    "                  f\"(MSE: {mse_loss.item():.4f}, Aux: {aux_loss.item():.4f})\")\n",
    "\n",
    "# Run the training test\n",
    "test_moe_training()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. . 课后作业\n",
    "- 把 expert 改成 swishGLU 版本的 FFN 专家\n",
    "    - 可以参考文章：[LLM activate function 激活函数的进化之路，从 ReLU，GELU 到 swishGLU 以及 手写 transformer decoder（CausalLM）](https://bruceyuan.com/llms-zero-to-hero/activate-function-from-relu-gelu-to-swishglu.html)\n",
    "- 把 MOE 应用到上一次的 build_nanoGPT 中，也就是替换掉原来的 FFN层，注意这里负载均衡 loss 要包含每一层的 MOE 的 router_logits\n",
    "    - 参考 GitHub 仓库， 【[LLMs-Zero-to-Hero](https://github.com/bbruceyuan/LLMs-Zero-to-Hero)】\n",
    "- 自己问一下 GPT topK 是怎么实现的反向传播，了解反向传播的梯度怎么流转的"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最后交个朋友，来群里吹水吧（🤣误）\n",
    "\n",
    "![chaofa-wechat-account:bbruceyuan](https://bruceyuan.com/llms-zero-to-hero/wechat-account-bbruceyuan.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
